from collections import OrderedDict
import os
from lib.sam2bed import *

class clust_align:
    def __init__(self,clus_obj,string):
        self.obj=clus_obj
        self.bedfile=string

class pairs:
	def __init__(self):
		self.p2={}
		self.st={}
	def set_dist(self,p,d,st):
		self.p2[p]=d
		self.st[p]=st

class sequence:
    def __init__(self, seq,freq,seq_id,pos=0):
        self.seq=seq
        self.freq=int(freq)
        self.len=len(seq)
        self.pos=pos
        self.id=seq_id

class position:
    def __init__(self,idl, chr,start,end,strand):
        self.idl=idl
        self.chr=chr
        self.start=int(start)
        self.end=int(end)
        self.strand=strand

class annotation:
    def __init__(self,db,name,strand,to5,to3):
        self.db=db
        self.name=name
        self.strand=strand
        self.to5=to5
        self.to3=to3

class dbannotation:
    def __init__(self,na):
        self.ann={}
    def adddbann(self,idl,ndba):
        self.ann[idl]=ndba

       
class cluster:
    def __init__(self, id):
        self.id=id
        self.loci={}
        self.locimax=0
        self.members=[]
        self.dbann={}
        self.idmembers={}
        self.loci2seq={}
        self.pairs={}
        self.dist_pair={}
        self.ref=0
        self.score=0
        self.showseq=""
    def set_ref(self,r):
    	self.ref=r

    def addpair(self,p1,p2,d,s):
		if not (self.pairs.has_key(p1)):
			self.pairs[p1]=pairs()
			self.dist_pair[p1]=1
		self.pairs[p1].set_dist(p2,d,s)
		self.dist_pair[p1]+=1
    
    def addidmember(self,id,idl):
        lenid=len(id)
        for c in id:
            self.idmembers[int(c)]=0
            if not self.loci2seq.has_key(idl):
                self.loci2seq[idl]=[]
            self.loci2seq[idl].append(int(c))
        if (lenid>self.locimax):
            self.locimax=lenid

    def addloci(self,idl,np):
        self.loci[idl]=np
        
    
    def addmember(self,ns):
        self.members.append(ns)
        

    def adddb(self,db,ndb):
        self.dbann[db]=ndb



class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'

class bedaligned:
	def __init__(self,l):
	    l=l.strip()
	    cols = l.split("\t")
	    self.chr=cols[0]
	    self.start=cols[1]
	    self.end=cols[2]
	    self.name=cols[3]
	    self.att=cols[4]
	    self.strand=cols[5]

class mergealigned:
	def __init__(self,l):
		self.chr=l[0]
		self.strand=l.strand
		self.start=l.start
		self.end=l.end
		self.names=l.name.split(";")
		self.loci=l.score.split(",")

def mergeclus(cur_clus_obj,clus_obj):
	#new_clus_obj=cur_clus_obj
    #print "in mergeclus function %s" % (len(cur_clus_obj))
	#list all c3newc old clusters id
    for idc in cur_clus_obj.keys():
        #print "cluster %s " % idc
		#retrieve old_clus_id at clus_obj
		
        for m in cur_clus_obj[idc].idmembers.keys():
            #print "member %s" % m          
            for id_old in clus_obj[m].members:

                cur_clus_obj[idc].addmember(id_old)


    #print "end"
	
    return cur_clus_obj

def get_distance(p1,p2):
	if p1.strand=="+":
		d = p1.start-p2.start
	else:
		d = (p1.end-(p1.end-p1.start)) - (p2.end-(p2.end-p2.start))
	return d

def get_ini_str(n):
	#get space values
	return "".join(" " for i in range(n))

def readbowtie():
    dict={}
    f = open("temp.map", 'r')
    for line in f:
        line = line.strip()
        cols = line.split("\t")
        numism=line.count(">")
        mism=""
        if numism >=1:
            mism=cols[7]
        
        result="%s:%s %s %s;" % (cols[2],cols[3],cols[1],mism)
        if dict.has_key(cols[0]):
            store=dict[cols[0]]+result
            dict[cols[0]] = store   
        else:
            dict[cols[0]] = result  
    f.close()
    return dict

def what_is(file):
    f = open(file,'r')
    cols= f.readline().split("\t")
    f.close()
    if (cols[5]=="+" or cols[5]=="-"):
        return "bed"
    else:
        return "sam"

def init_numlocidb(beds):
    dict={}
    for filebed in beds:
        db= os.path.basename(filebed)
        dict[db]=0;
    return dict

def calc_complexity(nl):
    ns=len(nl)
    total=0.0
    for l in nl:
        total+=1.0/l
        #print total
    total/=ns
    return (total)

def show_seq(clus_obj,seq_obj):
    for idc in clus_obj.keys():
        #print "cluster %s" % idc
        abs_pos_list={}
        seq_located={}
        if (len(clus_obj[idc].members)>1):
            nloci=[]
            templist = []
            listmembers=clus_obj[idc].pairs.keys()
            lenlistmembers=len(listmembers)
            i=1
            while lenlistmembers>len(templist):
                i+=1
                #print "iter %s" % i
                templist=[]
                #print "while start %s" % listmembers
                for m in listmembers:
                    if i==2:
                        numlociseq=len(seq_obj[m].loci)
                        nloci.append(numlociseq)
                    if not (abs_pos_list):
                            abs_pos_list[m]=0
                            seq_located[m]=0
                            #showseq=get_ini_str(90)+seqinfo[m].members[0].seq
                            #print "%s %s %s" % (showseq,0,m)
                    if abs_pos_list.has_key(m):
                            refpos=abs_pos_list[m]
                    else:
                        templist.append(m)
                    #if not seq_located.has_key(m):
                   
                    for p in clus_obj[idc].pairs[m].p2.keys():
                        
                        d=clus_obj[idc].pairs[m].p2[p]
                        if clus_obj[idc].pairs[m].st[p]=="-":
                            calc_d=refpos+(len(seq_obj[m].members[0].seq)-len(seq_obj[p].members[0].seq)+d)
                        else:
                            calc_d=refpos-d
                        if not abs_pos_list.has_key(p):
                            abs_pos_list[p]=20000
                        #print "%s %s calc %s abs %s d %s" % (m,p,calc_d,abs_pos_list[p],d)
                        if calc_d < abs_pos_list[p]:
                            seq_located[p]=0
                            abs_pos_list[p]=calc_d
                            #print "pairs %s dist %s strand %s" % (p,d,setclus[idc].pairs[m].st[p])
                            #print "%s %s" % (seqinfo[m].members[0].seq,seqinfo[p].members[0].seq)
                            #seq=seqinfo[p].members[0].seq
                            #n=90+abs_pos_list[p]
                            #showseq=get_ini_str(n)+seq
                            #print "%s ref %s d %s p %s id %s" % (showseq,refpos,d,setclus[idc].pairs[m].st[p],p)
                lenlistmembers=len(listmembers)
                listmembers=templist
            
            abs_pos_list_order = OrderedDict(sorted(abs_pos_list.items(), key=lambda x: x[1]))
            ini=-500
            for s in abs_pos_list_order.keys():
                if (ini==-500):
                    ini=abs_pos_list[s]
                seq=seq_obj[s].members[0].seq
                #n=100+abs_pos_list[s]
                n=abs(ini-abs_pos_list[s])
                showseq=get_ini_str(n)+seq
                #print "%s  id %s idini %s p %s" % (showseq,s,seq_obj[s].members[0].id,abs_pos_list[s]) 
                clus_obj[idc].showseq=clus_obj[idc].showseq+"\n" + showseq+"\t"+str(seq_obj[s].members[0].freq)
            clus_obj[idc].score=calc_complexity(nloci)
            #print "score cluster %s " % (score)
        else:
            #print "unique sequence %s" % setclus[idc].members[0].seq
            clus_obj[idc].showseq=clus_obj[idc].members[0].seq
            clus_obj[idc].score=1

    return clus_obj

def anncluster(c,clus_obj,db,type_ann):
    if type_ann=="gtf":
        id_sa=1
        id_ea=2
        id_sb=9
        id_eb=10
        id_id=3
        id_idl=4
        id_sta=5
        id_stb=12
        id_tag=8
    else:
        id_sa=1
        id_ea=2
        id_sb=7
        id_eb=8
        id_id=3
        id_idl=4
        id_sta=5
        id_stb=11
        id_tag=9


    for cols in c.features():
        id=int(cols[id_id])
        idl=int(cols[id_idl])
        if (clus_obj.has_key(id)):
            clus=clus_obj[id]
            str="-"
            sa=int(cols[id_sa])
            ea=int(cols[id_ea])
            sb=int(cols[id_sb])
            eb=int(cols[id_eb])
            if (cols[id_sta] in cols[id_stb]):
                str="+"
            if (cols[id_sta] in "+" and cols[id_stb] in "+"):
                lento5=sa-sb+1
                lento3=ea-eb+1
            if (cols[id_sta] in "+" and cols[id_stb] in "-"):
                lento5=ea-sb+1
                lento3=sa-eb+1
            if (cols[id_sta] in "-" and cols[id_stb] in "+"):
                lento5=sa-eb+1
                lento3=ea-sb+1
            if (cols[id_sta] in "-" and cols[id_stb] in "-"):
                lento3=sa-sb+1
                lento5=ea-eb+1
            else:
                lento5=sa-sb+1
                lento3=ea-eb+1
            #print "EXISTS clus %s with db DBA %s" % (cols[9],db)
            if (clus.dbann.has_key(db)):
                #print "NEW clus %s with db DBA %s" % (cols[9],db)
                ann=annotation(db,cols[id_tag],str,lento5,lento3)
                tdb=clus.dbann[db]
                tdb.adddbann(idl,ann)
                clus.adddb(db,tdb)
            else:
                #print "UPDATE clus %s with db DBA %s" % (cols[9],db)
                ann=annotation(db,cols[8],str,lento5,lento3)
                tdb=dbannotation(1)
                tdb.adddbann(idl,ann)
                clus.adddb(db,tdb)
            
            clus_obj[id]=clus
    return clus_obj

def parse_align_file(file,seq_l,format):
    f = open(file, 'r')
    loc_id=1
    clus_obj={}
    bedfile_clusters=""
    for line in f:

        if format=="sam":
            line=processSAM(line)
        if line:
            a=bedaligned(line)
            id_clus=seq_l[a.name].id
            if (clus_obj.has_key(id_clus)):
                clus=clus_obj[id_clus]
            else:       
                clus=cluster(id_clus)
                clus.addmember(seq_l[a.name])
            loc_id=loc_id+1
            newpos=position(loc_id,a.chr,a.start,a.end,a.strand)
            clus.addloci(loc_id,newpos)
            clus_obj[id_clus]=clus
            #print "%s\t%s\t%s\t%s\t%s\t%s\n" % (a.chr,a.start,a.end,id_clus,loc_id,a.strand)
            bedfile_clusters+="%s\t%s\t%s\t%s\t%s\t%s\n" % (a.chr,a.start,a.end,id_clus,loc_id,a.strand)
        
    f.close()
    return clust_align(clus_obj,bedfile_clusters)

def parse_fasta_file(file):
    f = open(file, 'r')
    name=""
    seq_l={}
    index=1
    for line in f:
        line=line.strip()

        if (">" in line):
            #counts=line.replace(">seq_","")
            counts = re.search("x([0-9]+)", line)
            #print line
            name=line.replace(">","")
        else:
            seq=line
            #add seq to repository
            index=index+1
            new_s=sequence(seq,counts.group(1),index)
            seq_l[name]=new_s

       
    f.close()
    return seq_l

def parse_merge_file(c,setclus):
    clus_id={}
    currentClus={}
    index=0
    lindex=0
    currentClustersBed=""
    eindex=0

    for line in c.features():
        a=mergealigned(line)
        exists=0
        lindex=lindex+1
        #print line
        for e in a.names:
            if (clus_id.has_key(e)):
                exists=1
                eindex=clus_id[e]
                break;
        if (exists==1):
            for e in a.names: 
                clus_id[e]=eindex
        else:
            index=index+1
            eindex=index
            for e in a.names: 
                clus_id[e]=eindex

        if ( not currentClus.has_key(eindex)):   
            currentClus[eindex]=cluster(eindex)

        #dev::distance from merge-start to seqs-start
        #print "%s %s %s %s %s %s" % (lindex,a.loci[0],a.names[0],a.chr,a.start,a.end)
        
        
        newpos=position(lindex,a.chr,a.start,a.end,a.strand)
        currentClus[eindex].addloci(lindex,newpos)
        currentClus[eindex].addidmember(a.names,lindex)
        #dev
        for e1 in range(len(a.names)):
            for e2 in range(len(a.names)):
                if (e1!=e2):
                    p1=setclus[int(a.names[e1])].loci[int(a.loci[e1])]
                    p2=setclus[int(a.names[e2])].loci[int(a.loci[e2])]
                    d = get_distance(p1,p2)
                    #print "#matchs%s %s " % (a.names[e1],a.names[e2])
                    if not currentClus[eindex].pairs.has_key(a.names[e1]):
                        currentClus[eindex].addpair(int(a.names[e1]),int(a.names[e2]),d,a.strand)
                        #print "#%s %s " % (a.names[e1],a.names[e2])
                    elif (currentClus[eindex].pairs.has_key(a.names[e1]) and (not(currentClus[eindex].pairs[a.names[e1]].p2.has_key(a.names[e2])))):
                        currentClus[eindex].addpair(int(a.names[e1]),int(a.names[e2]),d,a.strand)
                        #print "##%s %s " % (a.names[e1],a.names[e2])
                    elif (currentClus[eindex].pairs.has_key(a.names[e1]) and (currentClus[eindex].pairs[a.names[e1]].p2.has_key(a.names[e2]))):
                        if (d<currentClus[eindex].pairs[a.names[e1]].p2[a.names[e2]]):
                            currentClus[eindex].addpair(int(a.names[e1]),int(a.names[e2]),d,a.strand)
                            #print "##%s %s " % (a.names[e1],a.names[e2])

        #currentClustersBed=currentClustersBed + "%s\t%s\t%s\t%s\t%s\t%s\n " % (a.chr,a.start,a.end,eindex,lindex,a.strand)
    return currentClus


def reduceloci(clus_obj):
    for idc in clus_obj.keys():
        clus=clus_obj[idc]
        #print idc
        totalseq=len(clus.idmembers.keys())
        temploci2seq={}
        temploci={}

        for idl in clus.loci2seq.keys():
            locimembers=len(clus.loci2seq[idl])
            ratio=1.0*clus.locimax/totalseq
            #print "%s %s %s" % (idl,totalseq,locimembers)
            if (ratio>=0.9 and locimembers==clus.locimax):
                temploci2seq[idl]=clus.loci2seq[idl]
                temploci[idl]=clus.loci[idl]
                #print "selected"
        if len(temploci.keys())>0:
            clus.loci2seq=temploci2seq
            clus.loci=temploci
            
    return clus_obj

def generate_position_bed(clus_obj):
    bedaligned=""
    for idc in clus_obj.keys():
        clus=clus_obj[idc]
        for idl in clus.loci.keys():
            pos=clus.loci[idl]
            bedaligned+= "%s\t%s\t%s\t%s\t%s\t%s\n" % (pos.chr,pos.start,pos.end,idc,idl,pos.strand)

    return bedaligned